import psycopg2
from tabulate import tabulate
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--dbname", required=True, help="数据库名称")
args = parser.parse_args()

dbname = args.dbname
dbuser = "postgres"
dbpassword = "123456"
dbhost = "192.168.90.63"

# 连接数据库
conn = psycopg2.connect(database=dbname, user=dbuser, password=dbpassword, host=dbhost)

def list_table_names_and_comments_markdown():
    cursor = conn.cursor()

    cursor.execute("SELECT table_name, obj_description(pg_class.oid) AS comment FROM information_schema.tables JOIN pg_class ON pg_class.relname = information_schema.tables.table_name WHERE table_schema = 'public' order by table_name asc")
    results = cursor.fetchall()

    markdown_table = "| 表名 | 描述 |\n|----|----|\n"
    for table_name, comment in results:
        markdown_table += f"| {table_name} | {comment} |\n"

    cursor.close()

    return markdown_table + "\n\n"


#  约束生成
def list_table_constrain(table_name):
    cursor = conn.cursor()
    cursor.execute(f"""
        SELECT
    con.conname AS "约束名称",
    CASE con.contype
        WHEN 'p' THEN 'p(主键)'::TEXT
        WHEN 'f' THEN 'f(外键)'::TEXT
        WHEN 'u' THEN 'u(唯一约束)'::TEXT
        WHEN 'c' THEN 'c(检查约束)'::TEXT
        WHEN 'x' THEN 'x(排除约束)'::TEXT
        ELSE con.contype::TEXT
    END AS "约束类型",
    con.condeferrable AS "是否可延迟",
    con.condeferred AS "是否默认延迟",
    con.convalidated AS "是否已验证",
    a.attname AS "相关字段",
    con.confrelid::regclass AS "引用表",
    CASE con.confupdtype
        WHEN 'a' THEN 'a(无动作)'::TEXT
        WHEN 'r' THEN 'r(限制)'::TEXT
        WHEN 'c' THEN 'c(级联)'::TEXT
        WHEN 'n' THEN 'n(置空)'::TEXT
        WHEN 'd' THEN 'd(置默认)'::TEXT
        ELSE con.confupdtype::TEXT
    END AS "更新规则",
    CASE con.confdeltype
        WHEN 'a' THEN 'a(无动作)'::TEXT
        WHEN 'r' THEN 'r(限制)'::TEXT
        WHEN 'c' THEN 'c(级联)'::TEXT
        WHEN 'n' THEN 'n(置空)'::TEXT
        WHEN 'd' THEN 'd(置默认)'::TEXT
        ELSE con.confdeltype::TEXT
    END AS "删除规则",
    CASE con.confmatchtype
        WHEN 'f' THEN 'f(全匹配)'::TEXT
        WHEN 'p' THEN 'p(部分匹配)'::TEXT
        WHEN 's' THEN 's(简单匹配)'::TEXT
        ELSE con.confmatchtype::TEXT
    END AS "匹配类型"
FROM
    pg_constraint con
JOIN
    pg_class rel
    ON rel.oid = con.conrelid
JOIN
    pg_attribute a
    ON a.attnum = ANY(con.conkey)
    AND a.attrelid = con.conrelid
WHERE
    rel.relname = '{table_name}';
        """)

    # 获取查询结果
    rows = cursor.fetchall()

    # 检查是否有结果
    if not rows:
        cursor.close()
        return ""
    
    # 获取列名
    columns = [desc[0] for desc in cursor.description]
    
    # 关闭游标和连接
    cursor.close()

    # 生成 Markdown 表格
    markdown_table = tabulate(rows, headers=columns, tablefmt="github")
    
    return markdown_table
    

# 数据字典生成
def generate_data_dict():

    cursor = conn.cursor()
    # 获取所有表名
    cursor.execute("SELECT table_name, obj_description(pg_class.oid) AS comment FROM information_schema.tables JOIN pg_class ON pg_class.relname = information_schema.tables.table_name WHERE table_schema = 'public'")
    cursor.execute("""
        SELECT TABLE_NAME,
            OBJ_DESCRIPTION(PG_CLASS.OID) AS COMMENT
        FROM INFORMATION_SCHEMA.TABLES
        JOIN PG_CLASS ON PG_CLASS.RELNAME = INFORMATION_SCHEMA.TABLES.TABLE_NAME
        WHERE TABLE_SCHEMA = 'public'
        """)
    tables = cursor.fetchall()

    markdown_content = ""

    for table_name, comment in tables:
        # 获取表的列信息
        # cursor.execute(f"SELECT column_name, col_description(pg_attribute.attrelid, pg_attribute.attnum) AS column_comment, data_type, is_nullable, column_default FROM information_schema.columns JOIN pg_attribute ON pg_attribute.attname = information_schema.columns.column_name AND pg_attribute.attrelid = (SELECT oid FROM pg_class WHERE relname = information_schema.columns.table_name) WHERE table_name = '{table_name}'")
        cursor.execute(f"""
            SELECT COLUMN_NAME,
                COL_DESCRIPTION(PG_ATTRIBUTE.ATTRELID, PG_ATTRIBUTE.ATTNUM) AS COLUMN_COMMENT,
                DATA_TYPE || 
    CASE 
        WHEN DATA_TYPE IN ('character varying', 'varchar', 'character', 'char') 
             AND character_maximum_length IS NOT NULL THEN 
            '(' || character_maximum_length || ')' 
        ELSE 
            '' 
    END AS DATA_TYPE_WITH_LENGTH,
                IS_NULLABLE
            FROM INFORMATION_SCHEMA.COLUMNS
            JOIN PG_ATTRIBUTE ON PG_ATTRIBUTE.ATTNAME = INFORMATION_SCHEMA.COLUMNS.COLUMN_NAME
            AND PG_ATTRIBUTE.ATTRELID =
                (SELECT OID
                    FROM PG_CLASS
                    WHERE RELNAME = INFORMATION_SCHEMA.COLUMNS.TABLE_NAME)
            WHERE TABLE_NAME = '{table_name}'
            """)
        columns = cursor.fetchall()

        # 获取索引信息
        cursor.execute(f"SELECT indexname, indexdef FROM pg_indexes WHERE tablename = '{table_name}'")
        indexColumns = cursor.fetchall()

        # 获取约束信息
        # cursor.execute(f"SELECT conname FROM pg_constraint WHERE conrelid = (SELECT oid FROM pg_class WHERE relname = '{table_name}')")
        # constraints = [constraint[0] for constraint in cursor.fetchall()]

        constrain_content = list_table_constrain(table_name)

        # 生成 Markdown 表格
        # table_content = tabulate(columns, headers=["列名", "列描述", "数据类型", "是否可为空", "默认值"], tablefmt="github")
        table_content = tabulate(columns, headers=["列名", "列描述", "数据类型", "是否可为空"], tablefmt="github")
        markdown_content += f"## {table_name} - {comment}\n\n索引: \n"

        index_output = ""
        for indexname, indexdef in indexColumns:
          # 在这里对获取到的每个值进行处理或打印等操作
          index_output += "Index Name: " + indexname + "\n"
          index_output += "Index Definition: " + indexdef + "\n\n"
        markdown_content += index_output

        # markdown_content += f"约束: {', '.join(constraints)}\n\n{table_content}\n\n\n"
        markdown_content += "约束: \n" + constrain_content
        markdown_content += f"\n\n{table_content}\n\n\n"

    # 关闭连接
    cursor.close()

    return markdown_content

# 生成并打印数据字典
data_dict = list_table_names_and_comments_markdown() + generate_data_dict()

conn.close()

with open(f"{dbname}.md", 'w') as f:
    f.write(data_dict)